package com.castile.processors.standard.syslog;

import com.castile.ssl.context.CustomSSLContextService;
import org.apache.commons.lang3.StringUtils;
import org.apache.nifi.annotation.behavior.InputRequirement;
import org.apache.nifi.annotation.behavior.TriggerWhenEmpty;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.SeeAlso;
import org.apache.nifi.annotation.documentation.Tags;
import org.apache.nifi.annotation.lifecycle.OnScheduled;
import org.apache.nifi.annotation.lifecycle.OnStopped;
import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.components.PropertyValue;
import org.apache.nifi.components.ValidationContext;
import org.apache.nifi.components.ValidationResult;
import org.apache.nifi.event.transport.EventSender;
import org.apache.nifi.event.transport.configuration.LineEnding;
import org.apache.nifi.event.transport.configuration.TransportProtocol;
import org.apache.nifi.event.transport.netty.StringNettyEventSenderFactory;
import org.apache.nifi.flowfile.FlowFile;
import org.apache.nifi.processor.ProcessContext;
import org.apache.nifi.processor.ProcessSession;
import org.apache.nifi.processor.ProcessorInitializationContext;
import org.apache.nifi.processor.Relationship;
import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.processors.standard.AbstractSyslogProcessor;
import org.apache.nifi.processors.standard.ListenSyslog;
import org.apache.nifi.processors.standard.ParseSyslog;
import org.apache.nifi.processors.standard.PutSyslog;
import org.apache.nifi.syslog.parsers.SyslogParser;
import org.apache.nifi.util.StopWatch;

import javax.net.ssl.SSLContext;
import java.nio.charset.Charset;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static org.apache.nifi.processors.standard.PutSyslog.*;

/**
 * 自定义的syslog处理器
 *
 * @author castile
 * @date 2025-06-13 下午11:49
 */
@InputRequirement(InputRequirement.Requirement.INPUT_REQUIRED)
@TriggerWhenEmpty
@Tags({"syslog", "put", "udp", "tcp", "logs"})
@CapabilityDescription("Sends Syslog messages to a given host and port over TCP or UDP. Messages are constructed from the \"Message ___\" properties of the processor " +
        "which can use expression language to generate messages from incoming FlowFiles. The properties are used to construct messages of the form: " +
        "(<PRIORITY>)(VERSION )(TIMESTAMP) (HOSTNAME) (BODY) where version is optional.  The constructed messages are checked against regular expressions for " +
        "RFC5424 and RFC3164 formatted messages. The timestamp can be an RFC5424 timestamp with a format of \"yyyy-MM-dd'T'HH:mm:ss.SZ\" or \"yyyy-MM-dd'T'HH:mm:ss.S+hh:mm\", " +
        "or it can be an RFC3164 timestamp with a format of \"MMM d HH:mm:ss\". If a message is constructed that does not form a valid Syslog message according to the " +
        "above description, then it is routed to the invalid relationship. Valid messages are sent to the Syslog server and successes are routed to the success relationship, " +
        "failures routed to the failure relationship.")
@SeeAlso({ListenSyslog.class, ParseSyslog.class})
public class CustomPutSyslog extends AbstractSyslogProcessor {

    private Set<Relationship> relationships;
    private List<PropertyDescriptor> descriptors;

    private EventSender<String> eventSender;
    private String transitUri;

    public static final PropertyDescriptor SSL_CONTEXT_SERVICE = new PropertyDescriptor.Builder()
            .name("SSL Context Service")
            .description("The Controller Service to use in order to obtain an SSL Context. If this property is set, syslog " +
                    "messages will be sent over a secure connection.")
            .required(false)
            .identifiesControllerService(CustomSSLContextService.class)
            .dependsOn(PROTOCOL, TCP_VALUE)
            .build();
    @Override
    protected void init(final ProcessorInitializationContext context) {
        final List<PropertyDescriptor> descriptors = new ArrayList<>();
        descriptors.add(HOSTNAME);
        descriptors.add(PROTOCOL);
        descriptors.add(PORT);
        descriptors.add(MAX_SOCKET_SEND_BUFFER_SIZE);
        descriptors.add(SSL_CONTEXT_SERVICE);
        descriptors.add(IDLE_EXPIRATION);
        descriptors.add(TIMEOUT);
        descriptors.add(BATCH_SIZE);
        descriptors.add(CHARSET);
        descriptors.add(MSG_PRIORITY);
        descriptors.add(MSG_VERSION);
        descriptors.add(MSG_TIMESTAMP);
        descriptors.add(MSG_HOSTNAME);
        descriptors.add(MSG_BODY);
        this.descriptors = Collections.unmodifiableList(descriptors);

        final Set<Relationship> relationships = new HashSet<>();
        relationships.add(REL_SUCCESS);
        relationships.add(REL_FAILURE);
        relationships.add(REL_INVALID);
        this.relationships = Collections.unmodifiableSet(relationships);
    }

    @Override
    public final List<PropertyDescriptor> getSupportedPropertyDescriptors() {
        return descriptors;
    }

    @Override
    protected Collection<ValidationResult> customValidate(final ValidationContext context) {
        final Collection<ValidationResult> results = new ArrayList<>();

        final String protocol = context.getProperty(PROTOCOL).getValue();
        CustomSSLContextService controllerService = context.getProperty(SSL_CONTEXT_SERVICE).asControllerService(CustomSSLContextService.class);

        if (UDP_VALUE.getValue().equals(protocol) && controllerService != null) {
            results.add(new ValidationResult.Builder()
                    .explanation("SSL can not be used with UDP")
                    .valid(false).subject("SSL Context").build());
        }

        return results;
    }
    @OnScheduled
    public void onScheduled(final ProcessContext context) throws InterruptedException {
        eventSender = getEventSender(context);
        final String protocol = context.getProperty(PROTOCOL).getValue();
        final String hostname = context.getProperty(HOSTNAME).evaluateAttributeExpressions().getValue();
        final int port = context.getProperty(PORT).evaluateAttributeExpressions().asInteger();
        transitUri = String.format("%s://%s:%s", protocol, hostname, port);
    }

    @OnStopped
    public void onStopped() throws Exception {
        if (eventSender != null) {
            eventSender.close();
        }
    }
    protected EventSender<String> getEventSender(final ProcessContext context) {
        final TransportProtocol protocol = TransportProtocol.valueOf(context.getProperty(PROTOCOL).getValue());
        final String hostname = context.getProperty(HOSTNAME).evaluateAttributeExpressions().getValue();
        final int port = context.getProperty(PORT).evaluateAttributeExpressions().asInteger();
        final Charset charset = Charset.forName(context.getProperty(CHARSET).evaluateAttributeExpressions().getValue());

        final LineEnding lineEnding = TransportProtocol.TCP.equals(protocol) ? LineEnding.UNIX : LineEnding.NONE;
        final StringNettyEventSenderFactory factory = new StringNettyEventSenderFactory(getLogger(), hostname, port, protocol, charset, lineEnding);
        factory.setThreadNamePrefix(String.format("%s[%s]", PutSyslog.class.getSimpleName(), getIdentifier()));
        factory.setWorkerThreads(context.getMaxConcurrentTasks());
        factory.setMaxConnections(context.getMaxConcurrentTasks());

        final int timeout = context.getProperty(TIMEOUT).evaluateAttributeExpressions().asTimePeriod(TimeUnit.MILLISECONDS).intValue();
        factory.setTimeout(Duration.ofMillis(timeout));

        final PropertyValue sslContextServiceProperty = context.getProperty(SSL_CONTEXT_SERVICE);
        if (sslContextServiceProperty.isSet()) {
            final CustomSSLContextService sslContextService = sslContextServiceProperty.asControllerService(CustomSSLContextService.class);
            SSLContext sslContext = sslContextService.createContext();

            factory.setSslContext(sslContext);
        }

        return factory.getEventSender();
    }

    private String getSyslogMessage(final ProcessContext context, final FlowFile flowFile) {
        final String priority = context.getProperty(MSG_PRIORITY).evaluateAttributeExpressions(flowFile).getValue();
        final String version = context.getProperty(MSG_VERSION).evaluateAttributeExpressions(flowFile).getValue();
        final String timestamp = context.getProperty(MSG_TIMESTAMP).evaluateAttributeExpressions(flowFile).getValue();
        final String hostname = context.getProperty(MSG_HOSTNAME).evaluateAttributeExpressions(flowFile).getValue();
        final String body = context.getProperty(MSG_BODY).evaluateAttributeExpressions(flowFile).getValue();

        final StringBuilder messageBuilder = new StringBuilder();
        messageBuilder.append("<").append(priority).append(">");
        if (version != null) {
            messageBuilder.append(version).append(StringUtils.SPACE);
        }
        messageBuilder.append(timestamp).append(StringUtils.SPACE).append(hostname).append(StringUtils.SPACE).append(body);
        return messageBuilder.toString();
    }

    private boolean isValid(final String message) {
        for (final Pattern pattern : SyslogParser.MESSAGE_PATTERNS) {
            final Matcher matcher = pattern.matcher(message);
            if (matcher.matches()) {
                return true;
            }
        }
        return false;
    }

    @Override
    public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
        final int batchSize = context.getProperty(BATCH_SIZE).evaluateAttributeExpressions().asInteger();
        final List<FlowFile> flowFiles = session.get(batchSize);
        if (flowFiles.isEmpty()) {
            context.yield();
        } else {
            for (final FlowFile flowFile : flowFiles) {
                final StopWatch timer = new StopWatch(true);
                final String syslogMessage = getSyslogMessage(context, flowFile);
                if (isValid(syslogMessage)) {
                    try {
                        eventSender.sendEvent(syslogMessage);
                        timer.stop();

                        final long duration = timer.getDuration(TimeUnit.MILLISECONDS);
                        session.getProvenanceReporter().send(flowFile, transitUri, duration, true);

                        getLogger().debug("Send Completed {}", flowFile);
                        session.transfer(flowFile, REL_SUCCESS);
                    } catch (final Exception e) {
                        getLogger().error("Send Failed {}", flowFile, e);
                        session.transfer(flowFile, REL_FAILURE);
                    }
                } else {
                    getLogger().debug("Syslog Message Invalid {}", flowFile);
                    session.transfer(flowFile, REL_INVALID);
                }
            }
        }
    }
}
